import matplotlib.pyplot as plt
import torch
from dataSet import make_data1
from model import EIModel
from torch import nn
from torch import optim
import tqdm

# 数据的读取
X, Y = make_data1()

# 模型的构建
model = EIModel()

# 相关的配置
# 损失函数
loss_fn = nn.MSELoss()
# 优化器
opt = optim.SGD(model.parameters(), lr=0.0001)

# 模型的训练
train_loss = []
tqdm_range = tqdm.tqdm(range(500), total=500)
for epoch in tqdm_range:
    for x, y in zip(X, Y):
        y_pred = model(x)
        loss = loss_fn(y_pred, y)
        train_loss.append(loss.item())
        opt.zero_grad()
        loss.backward()
        opt.step()

print(train_loss)
# 模型的保存
torch.save(model.state_dict(), r'model.pth')
plt.plot(train_loss, label='train_loss')
plt.legend()
plt.show()
